#!/usr/bin/env python3
"""
输入提示模块

提供简单的接口来接收和传递点/框/掩码提示给SAM
"""

import torch
import numpy as np
from typing import Dict, List, Tuple, Optional, Union, Any
from dataclasses import dataclass

import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent.parent))
from utils.logger import get_logger


@dataclass
class PointPrompt:
    """点提示"""
    coords: torch.Tensor  # (N, 2) 坐标点
    labels: torch.Tensor  # (N,) 标签 (1=前景, 0=背景)


@dataclass
class BoxPrompt:
    """框提示"""
    boxes: torch.Tensor  # (N, 4) 边界框 [x1, y1, x2, y2]


@dataclass
class MaskPrompt:
    """掩码提示"""
    masks: torch.Tensor  # (N, H, W) 掩码


class InputPromptHandler:
    """输入提示处理器
    
    简单的接口来接收和格式化提示给SAM
    """
    
    def __init__(self, device: str = "cuda"):
        """
        初始化输入提示处理器
        
        Args:
            device: 设备
        """
        self.device = device
        self.logger = get_logger("InputPromptHandler")
    
    def process_point_prompt(self, 
                           coords: Union[torch.Tensor, np.ndarray, List],
                           labels: Optional[Union[torch.Tensor, np.ndarray, List]] = None) -> Dict[str, torch.Tensor]:
        """
        处理点提示
        
        Args:
            coords: 坐标点 (N, 2) 或 [(x1, y1), (x2, y2), ...]
            labels: 标签 (N,) 或 None（默认为前景）
            
        Returns:
            SAM格式的点提示
        """
        # 转换为tensor
        if isinstance(coords, (list, np.ndarray)):
            coords = torch.tensor(coords, dtype=torch.float32, device=self.device)
        else:
            coords = coords.to(self.device)
        
        # 处理标签
        if labels is None:
            labels = torch.ones(coords.shape[0], dtype=torch.int64, device=self.device)
        elif isinstance(labels, (list, np.ndarray)):
            labels = torch.tensor(labels, dtype=torch.int64, device=self.device)
        else:
            labels = labels.to(self.device)
        
        # 确保形状正确
        if len(coords.shape) == 1:
            coords = coords.unsqueeze(0)
        if len(labels.shape) == 0:
            labels = labels.unsqueeze(0)
        
        return {
            'point_coords': coords,
            'point_labels': labels
        }
    
    def process_box_prompt(self, 
                          boxes: Union[torch.Tensor, np.ndarray, List]) -> Dict[str, torch.Tensor]:
        """
        处理框提示
        
        Args:
            boxes: 边界框 (N, 4) 或 [[x1, y1, x2, y2], ...]
            
        Returns:
            SAM格式的框提示
        """
        # 转换为tensor
        if isinstance(boxes, (list, np.ndarray)):
            boxes = torch.tensor(boxes, dtype=torch.float32, device=self.device)
        else:
            boxes = boxes.to(self.device)
        
        # 确保形状正确
        if len(boxes.shape) == 1:
            boxes = boxes.unsqueeze(0)
        
        return {
            'boxes': boxes
        }
    
    def process_mask_prompt(self, 
                          masks: Union[torch.Tensor, np.ndarray]) -> Dict[str, torch.Tensor]:
        """
        处理掩码提示
        
        Args:
            masks: 掩码 (N, H, W) 或 (H, W)
            
        Returns:
            SAM格式的掩码提示
        """
        # 转换为tensor
        if isinstance(masks, np.ndarray):
            masks = torch.from_numpy(masks).float()
        
        masks = masks.to(self.device)
        
        # 确保形状正确
        if len(masks.shape) == 2:
            masks = masks.unsqueeze(0)
        
        return {
            'mask_inputs': masks
        }
    
    def format_for_sam(self, prompt: Dict[str, Any]) -> Dict[str, torch.Tensor]:
        """
        格式化提示为SAM输入格式
        
        Args:
            prompt: 包含 point_coords, point_labels, boxes, mask_inputs 的字典
            
        Returns:
            SAM格式的提示字典
        """
        sam_prompt = {}
        
        if 'point_coords' in prompt:
            sam_prompt['point_coords'] = prompt['point_coords']
        if 'point_labels' in prompt:
            sam_prompt['point_labels'] = prompt['point_labels']
        if 'boxes' in prompt:
            sam_prompt['boxes'] = prompt['boxes']
        if 'mask_inputs' in prompt:
            sam_prompt['mask_inputs'] = prompt['mask_inputs']
        
        return sam_prompt
